﻿using System;
using System.Collections.Generic;
using Microsoft.Xna.Framework;
using Microsoft.Xna.Framework.Content;
using Microsoft.Xna.Framework.Graphics;
using Phoenix.Xna.Framework.Components.Cameras;

namespace Phoenix.Xna.Framework.Graphics.Models
{
    public class InstancedModelPart
    {
        private const int MaxShaderMatrices = 60;
        private const int SizeOfVector4 = sizeof(float) * 4;
        private const int SizeOfMatrix = sizeof(float) * 16;
        [ContentSerializer]
        private int _indexCount = 0;
        [ContentSerializer]
        private int _vertexCount = 0;
        [ContentSerializer]
        private int _vertexStride = 0;
        [ContentSerializer]
        private VertexDeclaration _vertexDeclaration;
        [ContentSerializer]
        private VertexBuffer _vertexBuffer;
        [ContentSerializer]
        private IndexBuffer _indexBuffer;
        [ContentSerializer(SharedResource = true)]
        private Effect _effect = null;
        private bool _techniqueChanged;
        private GraphicsDevice _graphicsDevice;
        private int _maxInstances;
        private Matrix[] _tempMatrices = new Matrix[MaxShaderMatrices];
        private DynamicVertexBuffer _instanceDataStream;
        private VertexElement[] _originalVertexDeclaration;
        private bool _vertexDataIsReplicated;

        private InstancedModelPart()
        {
        }

        public void Draw(InstancingTechnique instancingTechnique, List<Matrix> instanceTransforms)
        {
            SetRenderStates(instancingTechnique);
            _effect.Begin();

            foreach (EffectPass pass in _effect.CurrentTechnique.Passes)
            {
                pass.Begin();

                switch (instancingTechnique)
                {
                    case InstancingTechnique.ShaderInstancing:
                        DrawShaderInstancing(instanceTransforms);
                        break;
                    case InstancingTechnique.HardwareInstancing:
                        DrawHardwareInstancing(instanceTransforms);
                        break;
                    case InstancingTechnique.NoInstancing:
                        DrawNoInstancing(instanceTransforms);
                        break;
                }

                pass.End();
            }

            _effect.End();
        }

        internal void Initialize(GraphicsDevice device)
        {
            _graphicsDevice = device;
            int indexOverflowLimit = ushort.MaxValue / _vertexCount;
            _maxInstances = Math.Min(indexOverflowLimit, MaxShaderMatrices);
            _originalVertexDeclaration = _vertexDeclaration.GetVertexElements();
        }

        internal void SetInstancingTechnique(InstancingTechnique instancingTechnique)
        {
            switch (instancingTechnique)
            {
                case InstancingTechnique.ShaderInstancing:
                    InitializeShaderInstancing();
                    break;
                case InstancingTechnique.HardwareInstancing:
                    InitializeHardwareInstancing();
                    break;
            }

            _techniqueChanged = true;
        }

        private void InitializeShaderInstancing()
        {
            if (!_vertexDataIsReplicated)
            {
                ReplicateVertexData();
                ReplicateIndexData();
                _vertexDataIsReplicated = true;
            }

            int instanceIndexOffset = _vertexStride - sizeof(float);
            byte usageIndex = 1;
            short stream = 0;

            VertexElement[] extraElements =
            {
                new VertexElement(
                    stream, 
                    (short)instanceIndexOffset,
                    VertexElementFormat.Single,
                    VertexElementMethod.Default,
                    VertexElementUsage.TextureCoordinate, 
                    usageIndex)
            };

            ExtendVertexDeclaration(extraElements);
        }

        private void InitializeHardwareInstancing()
        {
            VertexElement[] extraElements = new VertexElement[4];
            short offset = 0;
            byte usageIndex = 1;
            short stream = 1;

            for (int i = 0; i < extraElements.Length; i++)
            {
                extraElements[i] = new VertexElement(
                    stream,
                    offset,
                    VertexElementFormat.Vector4,
                    VertexElementMethod.Default,
                    VertexElementUsage.TextureCoordinate,
                    usageIndex);
                offset += SizeOfVector4;
                usageIndex++;
            }

            ExtendVertexDeclaration(extraElements);
        }

        private void ExtendVertexDeclaration(VertexElement[] extraElements)
        {
            _vertexDeclaration.Dispose();
            int length = _originalVertexDeclaration.Length + extraElements.Length;
            VertexElement[] elements = new VertexElement[length];
            _originalVertexDeclaration.CopyTo(elements, 0);
            extraElements.CopyTo(elements, _originalVertexDeclaration.Length);
            _vertexDeclaration = new VertexDeclaration(_graphicsDevice, elements);
        }

        private void ReplicateVertexData()
        {
            byte[] oldVertexData = new byte[_vertexCount * _vertexStride];
            _vertexBuffer.GetData(oldVertexData);
            _vertexBuffer.Dispose();
            int oldVertexStride = _vertexStride;
            _vertexStride += sizeof(float);
            byte[] newVertexData = new byte[_vertexCount * _vertexStride * _maxInstances];
            int outputPosition = 0;
            for (int instanceIndex = 0; instanceIndex < _maxInstances; instanceIndex++)
            {
                int sourcePosition = 0;
                byte[] instanceIndexBits = BitConverter.GetBytes((float)instanceIndex);

                for (int i = 0; i < _vertexCount; i++)
                {
                    Array.Copy(oldVertexData, sourcePosition, newVertexData, outputPosition, oldVertexStride);
                    outputPosition += oldVertexStride;
                    sourcePosition += oldVertexStride;
                    instanceIndexBits.CopyTo(newVertexData, outputPosition);
                    outputPosition += instanceIndexBits.Length;
                }
            }

            _vertexBuffer = new VertexBuffer(_graphicsDevice, newVertexData.Length, BufferUsage.None);
            _vertexBuffer.SetData(newVertexData);
        }

        private void ReplicateIndexData()
        {
            ushort[] oldIndices = new ushort[_indexCount];
            _indexBuffer.GetData(oldIndices);
            _indexBuffer.Dispose();
            ushort[] newIndices = new ushort[_indexCount * _maxInstances];
            int outputPosition = 0;

            for (int instanceIndex = 0; instanceIndex < _maxInstances; instanceIndex++)
            {
                int instanceOffset = instanceIndex * _vertexCount;

                for (int i = 0; i < _indexCount; i++)
                {
                    newIndices[outputPosition] = (ushort)(oldIndices[i] +
                                                          instanceOffset);

                    outputPosition++;
                }
            }

            _indexBuffer = new IndexBuffer(_graphicsDevice, sizeof(ushort) * newIndices.Length, BufferUsage.None, IndexElementSize.SixteenBits);
            _indexBuffer.SetData(newIndices);
        }

        private void SetRenderStates(InstancingTechnique instancingTechnique)
        {
            _graphicsDevice.VertexDeclaration = _vertexDeclaration;
            _graphicsDevice.Vertices[0].SetSource(_vertexBuffer, 0, _vertexStride);
            _graphicsDevice.Indices = _indexBuffer;

            if (_techniqueChanged)
            {
                string techniqueName = instancingTechnique.ToString();
                _effect.CurrentTechnique = _effect.Techniques[techniqueName];
                _techniqueChanged = false;
            }

            _effect.Parameters["View"].SetValue(Camera.ActiveCamera.View);
            _effect.Parameters["Projection"].SetValue(Camera.ActiveCamera.Projection);
        }

        private void DrawShaderInstancing(List<Matrix> instanceTransforms)
        {
            for (int i = 0; i < instanceTransforms.Count; i += _maxInstances)
            {
                int instanceCount = instanceTransforms.Count - i;

                if (instanceCount > _maxInstances)
                {
                    instanceCount = _maxInstances;
                }

                Array.Copy(instanceTransforms.ToArray(), i, _tempMatrices, 0, instanceCount);
                _effect.Parameters["InstanceTransforms"].SetValue(_tempMatrices);
                _effect.CommitChanges();
                _graphicsDevice.DrawIndexedPrimitives(PrimitiveType.TriangleList, 0, 0, instanceCount * _vertexCount, 0, instanceCount * _indexCount / 3);
            }
        }

        private void DrawHardwareInstancing(List<Matrix> instanceTransforms)
        {
            int instanceDataSize = SizeOfMatrix * instanceTransforms.Count;

            if ((_instanceDataStream == null) || (_instanceDataStream.SizeInBytes < instanceDataSize))
            {
                if (_instanceDataStream != null)
                {
                    _instanceDataStream.Dispose();
                }

                _instanceDataStream = new DynamicVertexBuffer(_graphicsDevice, instanceDataSize, BufferUsage.WriteOnly);
            }

            _instanceDataStream.SetData(instanceTransforms.ToArray(), 0, instanceTransforms.Count, SetDataOptions.Discard);
            VertexStreamCollection vertices = _graphicsDevice.Vertices;
            vertices[0].SetFrequencyOfIndexData(instanceTransforms.Count);
            vertices[1].SetSource(_instanceDataStream, 0, SizeOfMatrix);
            vertices[1].SetFrequencyOfInstanceData(1);
            _graphicsDevice.DrawIndexedPrimitives(PrimitiveType.TriangleList, 0, 0, _vertexCount, 0, _indexCount / 3);
            vertices[0].SetSource(null, 0, 0);
            vertices[1].SetSource(null, 0, 0);
        }

        private void DrawNoInstancing(List<Matrix> instanceTransforms)
        {
            EffectParameter transform = _effect.Parameters["NoInstancingTransform"];

            for (int i = 0; i < instanceTransforms.Count; i++)
            {
                transform.SetValue(instanceTransforms[i]);
                _effect.CommitChanges();
                _graphicsDevice.DrawIndexedPrimitives(PrimitiveType.TriangleList, 0, 0, _vertexCount, 0, _indexCount / 3);
            }
        }
    }
}
